from .base_strategy import BaseStrategy
import backtrader as bt
import indicator.MyTTIndicator as Mytt

###
# MACD 金叉死叉进行交易。
# 金叉做多，死叉开空
# 固定2%止损，固定10根k线后止盈
# #
class MyStrategy(BaseStrategy):
    params = (
        ('atr40_period', 40),
        ('atr130_period', 130),
        ('donchian_period', 130),
        ('atr_multiplier', 2),  # 止损倍数
    )
 
    def __init__(self):
        super().__init__()

        self.init_indicators()

        # 跟踪订单和交易状态
        self.order = None
        self.position_type = None  # 'long' or 'short'
    
    def init_indicators(self):
        """在此定义指标 (EMA/MACD/RSI等)"""
        # 计算ATR指标  并且不在图表上绘制   plot=False
        self.atr40 = bt.indicators.ATR(self.data, period=self.p.atr40_period,plot=False)
        self.atr130 = bt.indicators.ATR(self.data, period=self.p.atr130_period,plot=False)
        
        # 计算唐齐安通道
        self.donchian_high = bt.indicators.Highest(self.data.high, period=self.p.donchian_period,plot=False)
        self.donchian_low = bt.indicators.Lowest(self.data.low, period=self.p.donchian_period,plot=False)
        
        # 将通道绘制在 主图上
                # 添加绘图到主图
        # 第一个参数是数据源，后面是绘图线
        # 使用plotmaster=self.data 确保绘制在主图
        bt.indicators.Highest(self.data.high, period=self.p.donchian_period, 
                            plotmaster=self.data, subplot=False, plotname='up')
        bt.indicators.Lowest(self.data.low, period=self.p.donchian_period, 
                           plotmaster=self.data, subplot=False, plotname='low')
        # 可选：绘制中轨
        # bt.indicators.LineOver((self.donchian_high + self.donchian_low) / 2, 
        #                       plotmaster=self.data, subplot=False, plotname='唐齐安中轨')

        # 波动性条件：短期ATR > 长期ATR
        self.volatility_condition = self.atr40 > self.atr130
        # 添加波动性条件到副图
        # self.volatility_plot = bt.indicators.SignalLine(
        #     self.volatility_condition,
        #     plotname='bodong',
        #     plotmaster=None,  # 确保绘制在副图
        #     subplot=True,
        #     color='red'
        # )

        # 添加波动性条件到副图
        # 使用 Plot 机制手动绘制
 
        self.volatility_plot2 = Mytt.MyTT_ATR2(self.data,period=40,period2=130)

    def notify_order(self, order):
        super().notify_order(order)

        if order.status in [order.Submitted, order.Accepted]:
            # 订单已提交/接受 - 不做任何操作
            return
        
        if order.status in [order.Completed]:
            if order.isbuy():
                print(f'买入执行, 价格: {order.executed.price:.2f}, 成本: {order.executed.value:.2f}, 佣金: {order.executed.comm:.2f}')
                self.position_type = 'long'
            elif order.issell():
                print(f'卖出执行, 价格: {order.executed.price:.2f}, 成本: {order.executed.value:.2f}, 佣金: {order.executed.comm:.2f}')
                if self.position:  # 检查是否为平仓操作
                    self.position_type = None
                    
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            print('订单取消/保证金不足/拒绝')
            
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
            
        print(f'交易盈亏, 毛利: {trade.pnl:.2f}, 净利: {trade.pnlcomm:.2f}')

    def next(self):
        super().next()


            # 检查是否已有持仓
        if self.position:
            # 平仓逻辑
            if self.position_type == 'long':
                # 多头止损
                stop_price = self.position.price - self.atr_multiplier * self.atr40[0]
                if self.data.close[0] < stop_price:
                    self.close()  # 触发止损
                # 多头止盈：价格突破通道上轨
                elif self.data.close[0] > self.donchian_high[0]:
                    self.close()  # 平仓多头
                    
            elif self.position_type == 'short':
                # 空头止损
                stop_price = self.position.price + self.atr_multiplier * self.atr40[0]
                if self.data.close[0] > stop_price:
                    self.close()  # 触发止损
                # 空头止盈：价格跌破通道下轨
                elif self.data.close[0] < self.donchian_low[0]:
                    self.close()  # 平仓空头
        else:
            # 开仓逻辑 - 只有满足波动性条件时才交易
            if self.volatility_condition[0]:
                # 价格跌破唐齐安通道下轨，开多
                if self.data.close[0] < self.donchian_low[-1]:
                    self.order = self.buy()
                # 价格涨破唐齐安通道上轨，开空
                elif self.data.close[0] > self.donchian_high[-1]:
                    self.order = self.sell()